[Feat][KV offload][WIP] Separated process for CPU KV cache processing #22607
Conversation
Signed-off-by: ApostaC <[email protected]>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a separate process for CPU KV cache management, which is a significant architectural change. The implementation establishes the communication protocol and server structure using ZMQ. My review focuses on the completeness and robustness of this new component.
I've identified several areas where the implementation is incomplete, particularly concerning the 'onload' functionality, which is part of the feature's scope but not fully implemented in the protocol or server logic. Additionally, there are inconsistencies in logging practices, where print is used instead of the configured logger, which can hinder debugging and monitoring in production environments. These issues are critical to address to ensure the feature is robust and maintainable.
| ONLOAD_REQUEST = enum.auto() | ||
| ONLOAD_FINISHED = enum.auto() |
There was a problem hiding this comment.
The KVServerCmd enum includes ONLOAD_REQUEST and ONLOAD_FINISHED, but the corresponding message structures (e.g., KVServerOnloadRequest, KVServerOnloadFinished) are not defined. This incompleteness prevents the implementation of the onload functionality and will cause decode_payload to fail for these commands. Please define the necessary msgspec.Struct classes for onload operations and integrate them into the protocol, including the KVServerMsg union.
|
|
||
| def scheduler_process_response(socket: zmq.Socket, | ||
| finished_offloads: list[str], | ||
| finished_onloads: list[str]): |
There was a problem hiding this comment.
The finished_onloads parameter is unused because the function does not handle ONLOAD_FINISHED commands, despite the docstring indicating it should. This suggests the feature implementation is incomplete. Please add logic to process onload responses and populate this list, which will also require defining the ONLOAD_FINISHED message in protocol.py.
| def handle_onload_request(self, client_id, cmd, payload): | ||
| print("Received onload request from:", client_id) | ||
|
|
||
| # TODO: Do something here? |
There was a problem hiding this comment.
The handle_onload_request method is a stub that only prints to the console. This is a critical part of the KV cache offload/onload feature that remains unimplemented. Please implement the full logic for handling onload requests.
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: Yihua Cheng <[email protected]>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: Yihua Cheng <[email protected]>
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels. Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add 🚀 |
ApostaC
changed the title
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.Purpose
This PR is part of #22605
It introduces the structure of the CPU KV manager process and how the scheduler should talk with the process.
Test Plan
Test Result
(Optional) Documentation Update